package com.easy.mongodb.core.toolkit;

import com.easy.mongodb.common.constants.BaseMongoConstants;
import com.mongodb.client.MongoCursor;
import com.mongodb.client.MongoDatabase;
import org.bson.BsonValue;
import org.bson.Document;

import static com.easy.mongodb.common.constants.BaseMongoConstants.DEFAULT_MONGO_ID_NAME;
import static com.easy.mongodb.common.constants.BaseMongoConstants.LOCK_COLLECTION;

/**
 * 基于es写的轻量级分布式锁,仅供框架内部使用,可避免引入redis/zk等其它依赖
 * <p>
 * Copyright © 2022 xpc1024 All Rights Reserved
 **/
public class LockUtils {
    /**
     * id字段名
     */
    private final static String ID_FIELD = "_id";
    /**
     * 重试等待时间
     */
    private final static Integer WAIT_SECONDS = 60;
//    public String acquire(String key, long expiration) {
//        Query query = Query.query(Criteria.where("_id").is(key));
//        String token = this.generateToken();
//        Update update = new Update()
//                .setOnInsert("_id", key)
//                .setOnInsert("expireAt", System.currentTimeMillis() + expiration)
//                .setOnInsert("token", token);
//
//        FindAndModifyOptions options = new FindAndModifyOptions().upsert(true)
//                .returnNew(true);
//        LockDocument doc = mongoTemplate.findAndModify(query, update, options,
//                LockDocument.class);
//        boolean locked = doc.getToken() != null && doc.getToken().equals(token);
//
//        // 如果已过期
//        if (!locked && doc.getExpireAt() < System.currentTimeMillis()) {
//            DeleteResult deleted = this.mongoTemplate.remove(
//                    Query.query(Criteria.where("_id").is(key)
//                            .and("token").is(doc.getToken())
//                            .and("expireAt").is(doc.getExpireAt())),
//                    LockDocument.class);
//            if (deleted.getDeletedCount() >= 1) {
//                // 成功释放锁， 再次尝试获取锁
//                return this.acquire(key, expiration);
//            }
//        }
//
//        log.debug("Tried to acquire lock for key {} with token {} . Locked: {}",
//                key, token, locked);
//        return locked ? token : null;
//    }

    /**
     * 尝试获取es分布式锁
     *
     * @param client   MongoClient
     * @param idValue  id字段值实际未entityClass名,一个entity对应一把锁
     * @param maxRetry 最大重试次数
     * @return 是否获取成功
     */
    public static synchronized boolean tryLock(MongoDatabase client, String idValue, Integer maxRetry) {
        boolean existsIndex = false;

        MongoCursor<String> it = client.listCollectionNames().iterator();
        while (it.hasNext()) {
            String name = it.next();
            if (LOCK_COLLECTION.equalsIgnoreCase(name)) {
                existsIndex = true;
                break;
            }
        }
        if (!existsIndex) {
            client.createCollection(LOCK_COLLECTION);
        }
        if (maxRetry <= BaseMongoConstants.ZERO) {
            return Boolean.FALSE;
        }

        if (getCount(client, idValue) > BaseMongoConstants.ZERO) {
            try {
                Thread.sleep(WAIT_SECONDS / maxRetry);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            return tryLock(client, idValue, --maxRetry);
        } else {
            return createLock(client, idValue);
        }
    }

    /**
     * 创建锁
     *
     * @param client  MongoDatabase
     * @param idValue id字段值实际未entityClass名,一个entity对应一把锁
     * @return 是否创建成功
     */
    private static boolean createLock(MongoDatabase client, String idValue) {
        BsonValue insertCount;
        try {
            insertCount = client.getCollection(LOCK_COLLECTION).insertOne(new Document(DEFAULT_MONGO_ID_NAME, idValue)).getInsertedId();
            return idValue.equals(insertCount.asString().getValue());
        } catch (Exception e) {
            return false;
        }
    }

    /**
     * 释放锁
     *
     * @param client   MongoDatabase
     * @param idValue  id字段值实际未entityClass名,一个entity对应一把锁
     * @param maxRetry 最大重试次数
     * @return 是否释放成功
     */
    public synchronized static boolean release(MongoDatabase client, String idValue, Integer maxRetry) {
        Long deleteCount = -1L;
        if (maxRetry <= BaseMongoConstants.ZERO) {
            return Boolean.FALSE;
        }

        try {
            deleteCount = client.getCollection(LOCK_COLLECTION).deleteMany(new Document(DEFAULT_MONGO_ID_NAME, idValue)).getDeletedCount();
        } catch (Exception e) {
            e.printStackTrace();
            return retryRelease(client, idValue, --maxRetry);
        }
        if (deleteCount > 0L) {
            return Boolean.TRUE;
        } else {
            return retryRelease(client, idValue, maxRetry);
        }
    }

    /**
     * 重试释放
     *
     * @param client   MongoClient
     * @param idValue  id字段值实际未entityClass名,一个entity对应一把锁
     * @param maxRetry 最大重试次数
     * @return 是否重试成功
     */
    private static boolean retryRelease(MongoDatabase client, String idValue, Integer maxRetry) {
        try {
            Thread.sleep(WAIT_SECONDS / maxRetry);
        } catch (InterruptedException interruptedException) {
            interruptedException.printStackTrace();
        }
        return release(client, idValue, --maxRetry);
    }

    /**
     * 获取个数
     *
     * @param client  MongoClient
     * @param idValue id字段值实际未entityClass名,一个entity对应一把锁
     * @return 该id对应的锁的个数, 如果>0 说明已有锁,需重试获取,否则认为无锁
     */
    private static Integer getCount(MongoDatabase client, String idValue) {
        Long count = client.getCollection(LOCK_COLLECTION).countDocuments(new Document(DEFAULT_MONGO_ID_NAME, idValue));
        return count.intValue();
    }

}
